package com.wtp.wallpaper.utils.ai;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.gson.Gson;
import lombok.*;
import okhttp3.*;
import org.jetbrains.annotations.NotNull;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.io.FileOutputStream;
import java.io.IOException;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.text.SimpleDateFormat;
import java.util.*;

@AllArgsConstructor
@NoArgsConstructor
@Data
public class XFXHUtil {

    private String appid; //这里填写APPID
    private String apiSecret;
    private String apiKey;
    private String image_generate_url;
    private String image_understand_url;

    private static final Gson gson = new Gson();

    private final BigModelNew bigModelNew = new BigModelNew();

    public static final ObjectMapper mapper = new ObjectMapper();

    public byte[] generateImage(String prompt) throws Exception {
        String authUrl = getAuthUrl(image_generate_url, apiKey, apiSecret, "POST");
        String json = buildRequestBodyGenerateImage(prompt);
        String res = HttpUtil.doPostJson(authUrl, null, json);

        //System.out.println(res);
        JsonParse jsonParse = gson.fromJson(res, JsonParse.class);
        return Base64.getDecoder().decode(jsonParse.payload.choices.text.get(0).content);
    }

    public String[] understandImage(List<byte[]> bytes) throws Exception {
        return bigModelNew.understandImage(bytes);
    }


    private String buildRequestBodyGenerateImage(String prompt) {
        JSONObject root = new JSONObject();

        // 1. 构建 Header 节点
        JSONObject header = new JSONObject();
        header.put("app_id", appid);
        header.put("uid", generateShortUuid());
        root.put("header", header);

        // 2. 构建 Parameter 节点
        JSONObject parameter = new JSONObject();
        JSONObject chatParams = new JSONObject();
        chatParams.put("domain", "s291394db");
        chatParams.put("temperature", 0.5);
        chatParams.put("max_tokens", 4096);
        chatParams.put("width", 720);
        chatParams.put("height", 1280);
        parameter.put("chat", chatParams);
        root.put("parameter", parameter);

        // 3. 构建 Payload 节点
        JSONObject payload = new JSONObject();
        JSONObject message = new JSONObject();
        JSONArray textArray = new JSONArray();
        JSONObject textItem = new JSONObject();
        textItem.put("role", "user");
        textItem.put("content", prompt);
        textArray.add(textItem);
        message.put("text", textArray);
        payload.put("message", message);
        root.put("payload", payload);

        // 直接返回 JSON 字符串
        return root.toJSONString();
    }

    // 生成15位UUID工具方法
    private String generateShortUuid() {
        return UUID.randomUUID().toString().replace("-", "").substring(0, 15);
    }

    // 鉴权方法
    private String getAuthUrl(String hostUrl, String apiKey, String apiSecret, String requestMethod) throws Exception {
        //hostUrl = hostUrl.replace("ws://", "http://").replace("wss://", "https://");
        URL url = new URL(hostUrl);
        // 时间
        SimpleDateFormat format = new SimpleDateFormat("EEE, dd MMM yyyy HH:mm:ss z", Locale.US);
        format.setTimeZone(TimeZone.getTimeZone("GMT"));
        String date = format.format(new Date());
        // date="Thu, 12 Oct 2023 03:05:28 GMT";
        // 拼接
        String preStr = "host: " + url.getHost() + "\n" + "date: " + date + "\n" + requestMethod + " " + url.getPath() + " HTTP/1.1";
        // System.err.println(preStr);
        // SHA256加密
        Mac mac = Mac.getInstance("hmacsha256");
        SecretKeySpec spec = new SecretKeySpec(apiSecret.getBytes(StandardCharsets.UTF_8), "hmacsha256");
        mac.init(spec);

        byte[] hexDigits = mac.doFinal(preStr.getBytes(StandardCharsets.UTF_8));
        // Base64加密
        String sha = Base64.getEncoder().encodeToString(hexDigits);
        // System.err.println(sha);
        // 拼接
        String authorization = String.format("api_key=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"", apiKey, "hmac-sha256", "host date request-line", sha);
        // 拼接地址
        HttpUrl httpUrl = Objects.requireNonNull(HttpUrl.parse("https://" + url.getHost() + url.getPath())).newBuilder().//
                addQueryParameter("authorization", Base64.getEncoder().encodeToString(authorization.getBytes(StandardCharsets.UTF_8))).//
                addQueryParameter("date", date).//
                addQueryParameter("host", url.getHost()).//
                build();

        // System.err.println(httpUrl.toString());
        return httpUrl.toString();
    }

    class BigModelNew extends WebSocketListener {
        private String NewQuestion = "";
        private WebSocket webSocket;
        private StringBuffer totalAnswer;
        private boolean wsCloseFlag;
        private List<byte[]> bytes;

        public String[] understandImage(List<byte[]> bytes) throws Exception {

            String authUrl = getAuthUrl(image_understand_url, apiKey, apiSecret, "GET");
            //String json = buildRequestBodyUnderstandImage(bytes);
            String url = authUrl.replace("http://", "ws://").replace("https://", "wss://");
            OkHttpClient client = new OkHttpClient.Builder().build();
            Request request = new Request.Builder().url(url).build();

            wsCloseFlag = false;
            totalAnswer = new StringBuffer();
            webSocket = client.newWebSocket(request, this);
            webSocket.send(buildRequestBodyUnderstandImage(bytes));
            while (!wsCloseFlag) {
                try {
                    Thread.sleep(200);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
            }
            return totalAnswer.toString().split("#");
        }

        private String buildRequestBodyUnderstandImage(List<byte[]> bytes) {

            JSONObject requestJson = new JSONObject();

            JSONObject header = new JSONObject();  // header参数
            header.put("app_id", appid);
            header.put("uid", UUID.randomUUID().toString().substring(0, 10));

            JSONObject parameter = new JSONObject(); // parameter参数
            JSONObject chat = new JSONObject();
            chat.put("domain", "imagev3");
            chat.put("temperature", 0.5);
            chat.put("max_tokens", 4096);
            chat.put("auditing", "default");
            parameter.put("chat", chat);

            JSONObject payload = new JSONObject(); // payload参数
            JSONObject message = new JSONObject();
            JSONArray text = new JSONArray();

            RoleContent roleContent;
            for (byte[] aByte : bytes) {
                roleContent = new RoleContent();
                roleContent.role = "user";
                roleContent.content = Base64.getEncoder().encodeToString(aByte);
                roleContent.content_type = "image";
                text.add(JSON.toJSON(roleContent));
            }

            // 添加对图片提出要求的信息
            roleContent = new RoleContent();
            roleContent.role = "user";
            roleContent.content = "请对以上图片生成各自的描述,不同图片描述之间使用#号分隔,返回一个总段落。只描述主要特征，不超过60字";
            roleContent.content_type = "text";
            text.add(JSON.toJSON(roleContent));

            message.put("text", text);
            payload.put("message", message);

            requestJson.put("header", header);
            requestJson.put("parameter", parameter);
            requestJson.put("payload", payload);

            // 直接返回 JSON 字符串
            return requestJson.toJSONString();
        }

/*        @Override
        public void onOpen(@NotNull WebSocket webSocket, @NotNull Response response) {
            super.onOpen(webSocket, response);
        }*/

        @Override
        public void onMessage(@NotNull WebSocket webSocket, @NotNull String text) {
            // System.out.println(userId + "用来区分那个用户的结果" + text);
            JsonParse myJsonParse = gson.fromJson(text, JsonParse.class);
            if (myJsonParse.header.code != 0) {
                System.out.println("发生错误，错误码为：" + myJsonParse.header.code);
                System.out.println("本次请求的sid为：" + myJsonParse.header.sid);
                System.out.println("本次请求的message为：" + myJsonParse.header.message);
                webSocket.close(1000, "");
                wsCloseFlag = true;
                totalAnswer.append("暂无描述");
                return;
            }
            List<Text> textList = myJsonParse.payload.choices.text;
            for (Text temp : textList) {
                //System.out.print(temp.content);
                totalAnswer.append(temp.content);
            }
            if (myJsonParse.header.status == 2) {
                wsCloseFlag = true;
            }
        }

        @Override
        public void onFailure(@NotNull WebSocket webSocket, @NotNull Throwable t, Response response) {
            super.onFailure(webSocket, t, response);
            try {
                if (null != response) {
                    int code = response.code();
                    System.out.println("onFailure code:" + code);
                    System.out.println("onFailure body:" + response.body().string());
                    if (101 != code) {
                        System.out.println("connection failed");
                        System.exit(0);
                    }
                    webSocket.close(1000, "");
                    wsCloseFlag = true;
                }
            } catch (IOException e) {
                // TODO Auto-generated catch block
                e.printStackTrace();
            }
        }

    }

    @Data
    static class RoleContent {
        String role;
        String content;
        String content_type;
    }

    static class JsonParse {
        Header header;
        Payload payload;
    }

    static class Header {
        int code;
        int status;
        String message;
        String sid;
    }

    static class Payload {
        Choices choices;
    }

    static class Choices {
        List<Text> text;
    }

    static class Text {
        String content;
    }
}
